package cn.kimmking.research.qedis.core.opr;

import cn.kimmking.research.qedis.core.AbstractOperator;
import lombok.AllArgsConstructor;
import lombok.Data;

import java.util.Arrays;
import java.util.LinkedHashSet;
import java.util.Objects;

/**
 * Zset operators.
 *
 * @Author : kimmking(kimmking@apache.org)
 * @create 2024/6/11 下午8:05
 */
public class ZsetOperator extends AbstractOperator {

    public int zadd(String key, String[] values, double[] scores) {
        CacheEntry<?> entry = getOrInitCacheEntry(key);
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) entry.getValue();
        for(int i = 0; i < values.length; i ++) {
            exist.add(new ZsetEntry(values[i], scores[i]));
        }
        return values.length;
    }

    public int zcard(String key) {
        if(checkInvalid(key)) return 0;
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) getEntryValue(key);
        return exist.size();
    }

    public int zcount(String key, double min, double max) {
        if(checkInvalid(key)) return 0;
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) getEntryValue(key);
        return (int) exist.stream().filter(e -> e.getScore() >= min && e.getScore() <= max).count();
    }

    public Double zscore(String key, String value) {
        if(checkInvalid(key)) return null;
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) getEntryValue(key);
        return exist.stream().filter(e -> e.getValue().equals(value))
                .map(ZsetEntry::getScore).findFirst().orElse(null);
    }

    public Integer zrank(String key, String value) {
        if(checkInvalid(key)) return null;
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) getEntryValue(key);
        Double zscore = zscore(key, value);
        if(zscore == null) return null;
        return (int) exist.stream().filter(x -> x.getScore() < zscore).count();
    }

    public int zrem(String key, String... values) {
        if(checkInvalid(key)) return 0;
        LinkedHashSet<ZsetEntry> exist = (LinkedHashSet<ZsetEntry>) getEntryValue(key);
        return (int) Arrays.stream(values)
                .map(x->exist.removeIf(y->y.getValue().equals(x)))
                .filter(Objects::nonNull).count();
    }

    private CacheEntry<?> getOrInitCacheEntry(String key) {
        CacheEntry<?> entry = getEntry(key);
        if(entry == null) {
            synchronized (getMap()) {
                if((entry = getEntry(key)) == null) {
                    entry = new CacheEntry<>(new LinkedHashSet<String>(), System.currentTimeMillis(), -1000);
                    putEntry(key, entry);
                }
            }
        }
        return entry;
    }

    @Data
    @AllArgsConstructor
    public static class ZsetEntry {
        private String value;
        private Double score;
    }
}
